import os
import cv2
import torch
import torch.nn as nn
from tqdm import tqdm
from utils.config import *
from core.gc_net import GCNET
from dataloader.data_loader import test_loader

# 测试
def test(model, test_loader, criterion):
    step = 0
    num_batches = len(test_loader)
    avg_error_3px = 0.0

    for batch in tqdm(test_loader): # 创建进度条
        # 从batch中取样本
        left_img = batch['left'].to(device)
        right_img = batch['right'].to(device)
        img_name = batch['name'][0]

        # 前向传播
        with torch.no_grad():
            disp = model(left_img, right_img)

        # 截断张量到[0, 255]，并转为灰度图像
        disp_img = torch.clamp(disp, min=0, max=255)[0, :].cpu().numpy()
        if 'usvinland' in dataset: disp_img = disp_img * 255.0 / 50.0
        # 映射到伪彩色，近红远蓝
        disp_pseudo = cv2.applyColorMap(cv2.convertScaleAbs(disp_img, alpha=5), cv2.COLORMAP_JET)

        # 保存预测视差和伪彩色图像
        out_path = os.path.join('./out', dataset)
        if not os.path.exists(os.path.join(out_path, 'pseudo')):
            os.makedirs(os.path.join(out_path, 'pseudo'))
        cv2.imwrite(os.path.join(out_path, img_name), disp_img)
        cv2.imwrite(os.path.join(out_path, 'pseudo', img_name), disp_pseudo)

        # 若不是KITTI数据集，计算误差
        if dataset != 'kitti2015':
            target_disp = batch['disp'].to(device)

            # 掩码
            mask = (target_disp < max_disp) & (target_disp > 0)
            mask = mask.detach_()

            # 计算损失
            loss = criterion(disp[mask], target_disp[mask]) # 只对mask=True的值位置计算损失

            # 计算batch的误差
            epe = torch.abs(disp[mask] - target_disp[mask]) # 预测视差和真实视差的差值
            error_mask = (epe >= 3.0) & (epe >= (target_disp[mask] * 0.05)) # 掩码，误差绝对值>3px且相对误差>5%
            error_3px = torch.sum(error_mask).item() / torch.numel(disp[mask]) * 100 # 错误像素占总像素的比例
            avg_error_3px += error_3px

            print('step: {:06} | total loss: {:8.5f} | error_3px: {:7.5}%'.format(step, loss.item(), error_3px))

        step += 1

    print('disp in out, pseudo in out/pseudo')

    if dataset != 'kitti2015':
        avg_error_3px = avg_error_3px / num_batches

        print('3px-error: {:7.5}%'.format(avg_error_3px))

# 测试最优模型，生成视差图
def main():
    assert BATCH_SIZE == 1, "Batch size should equal to 1"

    # 打印数据集名称
    print('dataset:', dataset)

    # 创建模型实例
    model = GCNET(height, width, max_disp)
    model = nn.DataParallel(model, device_ids=[0]).to(device)
    torch.backends.cudnn.benchmark = True

    # 加载模型
    if 'kitti' in dataset: best_model_path = './save/' + dataset + '/kitti.pth' # 模型路径
    if 'usvinland' in dataset: best_model_path = './save/' + dataset + '/usvinland.pth'
    if os.path.exists(best_model_path):
        state = torch.load(best_model_path)
        model.load_state_dict(state['state_dict'])

        # 损失函数
        criterion = torch.nn.L1Loss().to(device)

        # 测试并写入相应文件夹
        model.cuda()
        model.eval()
        test(model, test_loader, criterion)
    else:
        print('best model unexist')

if __name__ == '__main__':
    main()
